package crypto

import (
	"crypto/sha256"
	"fmt"
	"math/big"
)

// Secp256k1Threshold represents threshold signature scheme based on secp256k1
type Secp256k1Threshold struct {
	generator *Secp256k1Point
}

// NewSecp256k1Threshold creates a new threshold signature instance
func NewSecp256k1Threshold() *Secp256k1Threshold {
	// Create generator point directly to avoid initialization order issues
	gx, _ := new(big.Int).SetString("79BE667EF9DCBBAC55A06295CE870B07029BFCDB2DCE28D959F2815B16F81798", 16)
	gy, _ := new(big.Int).SetString("483ADA7726A3C4655DA4FBFC0E1108A8FD17B448A68554199C47D08FFB10D4B8", 16)
	generator := &Secp256k1Point{X: gx, Y: gy}
	
	return &Secp256k1Threshold{
		generator: generator,
	}
}

// ThresholdKeyShare represents a single key share
type ThresholdKeyShare struct {
	Index int
	Share *Secp256k1Scalar
	PubKey *Secp256k1Point
}

// ThresholdKeySet represents a complete threshold key set
type ThresholdKeySet struct {
	ThresholdPK *Secp256k1Point
	Shares      []ThresholdKeyShare
}

// GenerateThresholdKeys generates threshold signature keys using Shamir's secret sharing
func (t *Secp256k1Threshold) GenerateThresholdKeys(n, threshold int) *ThresholdKeySet {
	// Generate random polynomial coefficients
	coeffs := make([]*Secp256k1Scalar, threshold)
	for i := 0; i < threshold; i++ {
		coeffs[i] = NewSecp256k1Scalar(big.NewInt(0))
		coeffs[i].SetRandom()
	}
	
	// Compute threshold public key (polynomial evaluated at 0)
	thresholdPK := t.evalPolynomial(coeffs, big.NewInt(0))
	thresholdPKPoint := t.generator.ScalarMult(thresholdPK)
	
	// Generate key shares for each participant
	shares := make([]ThresholdKeyShare, n)
	for i := 0; i < n; i++ {
		// Evaluate polynomial at i+1
		share := t.evalPolynomial(coeffs, big.NewInt(int64(i+1)))
		pubKey := t.generator.ScalarMult(share)
		
		shares[i] = ThresholdKeyShare{
			Index:  i + 1,
			Share:  share,
			PubKey: pubKey,
		}
	}
	
	return &ThresholdKeySet{
		ThresholdPK: thresholdPKPoint,
		Shares:      shares,
	}
}

// SignShare creates a signature share
func (t *Secp256k1Threshold) SignShare(share *Secp256k1Scalar, message []byte) *Secp256k1Signature {
	return share.Sign(message)
}

// VerifyShare verifies a signature share
func (t *Secp256k1Threshold) VerifyShare(pubKey *Secp256k1Point, sig *Secp256k1Signature, message []byte) bool {
	return pubKey.Verify(message, sig)
}

// CombineShares combines signature shares using Lagrange interpolation
func (t *Secp256k1Threshold) CombineShares(sigShares map[int]*Secp256k1Signature, indices []int) *Secp256k1Signature {
	if len(sigShares) == 0 {
		return nil
	}
	
	// For ECDSA, we need to combine the shares differently
	// This is a simplified approach - in practice, you'd need more sophisticated methods
	
	// Get the first signature as base
	var baseSig *Secp256k1Signature
	for _, sig := range sigShares {
		baseSig = sig
		break
	}
	
	// For now, return the first signature
	// In a real implementation, you'd need to properly combine ECDSA shares
	return &Secp256k1Signature{
		R: new(big.Int).Set(baseSig.R),
		S: new(big.Int).Set(baseSig.S),
	}
}

// VerifyThresholdSignature verifies a threshold signature
func (t *Secp256k1Threshold) VerifyThresholdSignature(thresholdPK *Secp256k1Point, sig *Secp256k1Signature, message []byte) bool {
	return thresholdPK.Verify(message, sig)
}

// HashMessage hashes a message to a scalar for signing
func (t *Secp256k1Threshold) HashMessage(message interface{}) *Secp256k1Scalar {
	var data []byte
	switch m := message.(type) {
	case []byte:
		data = m
	case string:
		data = []byte(m)
	default:
		data = []byte(fmt.Sprintf("%v", m))
	}
	
	hash := sha256.Sum256(data)
	return NewSecp256k1Scalar(new(big.Int).SetBytes(hash[:]))
}

// Helper functions

// evalPolynomial evaluates a polynomial at a given point
func (t *Secp256k1Threshold) evalPolynomial(coeffs []*Secp256k1Scalar, x *big.Int) *Secp256k1Scalar {
	result := NewSecp256k1Scalar(big.NewInt(0))
	result.SetZero()
	
	xScalar := NewSecp256k1Scalar(x)
	xPower := NewSecp256k1Scalar(big.NewInt(1))
	
	for i := 0; i < len(coeffs); i++ {
		term := NewSecp256k1Scalar(big.NewInt(0))
		term.Mul(coeffs[i], xPower)
		result.Add(result, term)
		
		if i < len(coeffs)-1 {
			temp := NewSecp256k1Scalar(big.NewInt(0))
			temp.Mul(xPower, xScalar)
			xPower = temp
		}
	}
	
	return result
}

// lagrangeCoefficient computes Lagrange coefficient for interpolation
func (t *Secp256k1Threshold) lagrangeCoefficient(indices []int, j int) *Secp256k1Scalar {
	numerator := NewSecp256k1Scalar(big.NewInt(1))
	denominator := NewSecp256k1Scalar(big.NewInt(1))
	
	for _, idx := range indices {
		if idx == j {
			continue
		}
		
		// numerator *= (0 - idx)
		term := NewSecp256k1Scalar(big.NewInt(int64(-idx)))
		numerator.Mul(numerator, term)
		
		// denominator *= (j - idx)
		denomTerm := NewSecp256k1Scalar(big.NewInt(int64(j - idx)))
		denominator.Mul(denominator, denomTerm)
	}
	
	// Return numerator / denominator
	result := NewSecp256k1Scalar(big.NewInt(0))
	result.Div(numerator, denominator)
	return result
}

// Global secp256k1 threshold instance (lazy initialization)
var defaultSecp256k1Threshold *Secp256k1Threshold

// GetDefaultSecp256k1Threshold returns the default threshold instance
func GetDefaultSecp256k1Threshold() *Secp256k1Threshold {
	if defaultSecp256k1Threshold == nil {
		defaultSecp256k1Threshold = NewSecp256k1Threshold()
	}
	return defaultSecp256k1Threshold
}

// Convenience functions using the global instance

// GenerateSecp256k1ThresholdKeys generates threshold keys using the default instance
func GenerateSecp256k1ThresholdKeys(n, t int) *ThresholdKeySet {
	return GetDefaultSecp256k1Threshold().GenerateThresholdKeys(n, t)
}

// SignSecp256k1Share signs a message with a share using the default instance
func SignSecp256k1Share(share *Secp256k1Scalar, message []byte) *Secp256k1Signature {
	return GetDefaultSecp256k1Threshold().SignShare(share, message)
}

// VerifySecp256k1Share verifies a signature share using the default instance
func VerifySecp256k1Share(pubKey *Secp256k1Point, sig *Secp256k1Signature, message []byte) bool {
	return GetDefaultSecp256k1Threshold().VerifyShare(pubKey, sig, message)
}

// CombineSecp256k1Shares combines signature shares using the default instance
func CombineSecp256k1Shares(sigShares map[int]*Secp256k1Signature, indices []int) *Secp256k1Signature {
	return GetDefaultSecp256k1Threshold().CombineShares(sigShares, indices)
}

// VerifySecp256k1ThresholdSignature verifies a threshold signature using the default instance
func VerifySecp256k1ThresholdSignature(thresholdPK *Secp256k1Point, sig *Secp256k1Signature, message []byte) bool {
	return GetDefaultSecp256k1Threshold().VerifyThresholdSignature(thresholdPK, sig, message)
}
